package top.continew.ai.config;

import com.alibaba.cloud.ai.memory.redis.JedisRedisChatMemoryRepository;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import top.continew.ai.advisor.LogAdvisor;
import top.continew.ai.advisor.ReReadingAdvisor;
import top.continew.ai.dto.AiProperties;

import java.util.Arrays;
import java.util.List;

/**
 * 配置类
 *
 * @author: lonni
 * @CreateTime: 2025-08-20  21:44
 * @Description:
 */
@Configuration
@EnableConfigurationProperties(AiProperties.class)
@Slf4j
public class AiConfiguration {

    private final AiProperties aiProperties;

    public AiConfiguration(AiProperties aiProperties) {
        this.aiProperties = aiProperties;
    }


//    @Bean
//    public VectorStore redisVectorStore() {
//        return RedisVectorStore.builder(jedisPooled, openAiEmbeddingModel)
//                .indexName("custom-index")                // 可选：默认 spring-ai-index
//                .prefix("custom-prefix")                  // 可选：默认 embedding:
//                .metadataFields(                         // 定义元数据字段
//                        RedisVectorStore.MetadataField.tag("country"),
//                        RedisVectorStore.MetadataField.numeric("year"))
//                .initializeSchema(true)                   // 是否初始化索引
//                .batchingStrategy(new TokenCountBatchingStrategy()) // 可选：默认 TokenCountBatchingStrategy
//                .build();
//
//    }


    /**
     * 增加内存数据库
     *
     * @param embeddingModel
     * @return
     */
    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        SimpleVectorStore.SimpleVectorStoreBuilder builder = SimpleVectorStore.builder(embeddingModel);
        return builder.build();

    }


    /**
     * 重新设置查询为空时返回的提示
     *
     * @return
     */
    @Bean
    public ContextualQueryAugmenter contextualQueryAugmenter() {
        ContextualQueryAugmenter build = ContextualQueryAugmenter.builder()
                .allowEmptyContext(false)
                .emptyContextPromptTemplate(
                        PromptTemplate.builder()
                                .template(aiProperties.getEmptyPrompt())
                                .build())
                .build();
        return build;

    }


    /**
     * 配置查询重写器
     * 用于优化用户查询，提高检索效果
     */
    @Bean
    public RewriteQueryTransformer rewriteQueryTransformer(ChatClient.Builder builder) {
        return RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .promptTemplate(
                        PromptTemplate.builder()
                                .template(aiProperties.getRewritePrompt())
                                .build()
                )
                .build();


    }

    /**
     * 利用大型语言模型将查询转化为多个语义各异的变体，以涵盖不同的视角，
     * 这对于获取更多背景信息以及提高找到相关结果的可能性非常有用
     *
     * @param builder
     * @return
     */
//    @Bean
//    public MultiQueryExpander multiQueryExpander(ChatClient.Builder builder) {
//        return MultiQueryExpander.builder()
//                .chatClientBuilder(builder)
//                .promptTemplate(
//                        PromptTemplate.builder()
//                                .template(aiProperties.getSystemPrompt())
//                                .build()
//                )
//                .build();
//
//    }


    /**
     * 增强QuestionAnswerAdvisor
     * 1）Pre-Retrieval ：负责对用户查询进行处理，以获得最理想的检索结果
     * 2）Retrieval ：负责查询诸如向量存储这样的数据系统，并获取最相关的文档
     * 3）Post-Retrieval ：负责对检索到的文档进行处理，以获得最佳的生成结果
     * 4）Generation ：负责根据用户查询以及检索到的文档生成最终的响应
     *
     * @param vectorStore
     * @param rewriteQueryTransformer
     * @param contextualQueryAugmenter
     * @return
     */
    public RetrievalAugmentationAdvisor retrievalAugmentationAdvisor(
            VectorStore vectorStore,
            RewriteQueryTransformer rewriteQueryTransformer,
            ContextualQueryAugmenter contextualQueryAugmenter

    ) {

        return RetrievalAugmentationAdvisor.builder()
                //配置基于向量库的查询
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(vectorStore)
                        .build())
                //配置空查询时返回的提示
                .queryAugmenter(contextualQueryAugmenter)
                //重写查询
                .queryTransformers(rewriteQueryTransformer)
                //配置后置监控
                .documentPostProcessors(((query, documents) -> {
                    log.info("Original query:{} ", query.text());
                    log.info("Retrieved documents: {}", documents.size());
                    return documents;
                }))


                .build();

    }


    /**
     * AI 基础对话（支持多轮对话记忆）
     * 基础实例 ,其他实例可在此基础上修改
     * @param builder
     * @return
     */
    @Bean("chatClient")
    public ChatClient chatClient(ChatClient.Builder builder,
                                 @Autowired JedisRedisChatMemoryRepository jedisRedisChatMemoryRepository
    ) {
        return builder
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(
                                        MessageWindowChatMemory.builder()
                                                .chatMemoryRepository(jedisRedisChatMemoryRepository)
                                                //设置最大存储20条
                                                .maxMessages(20)
                                                .build())
                                .build(),
                        //日志记录
                        new LogAdvisor(),
                        //重读功能
                        new ReReadingAdvisor()

                )

                .build();
    }
















}
